import torch
from torchvision import transforms
from PIL import Image

# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 类别名称
class_names = ['hard', 'soft']  # 请确保与训练时的顺序一致

# 加载模型
model = torch.load('package_classifier_full.pth', map_location=device)
model.eval()

# 预处理变换
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def predict_image(image):
    """预测单张图像的类别"""
    image_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(image_tensor)
        _, preds = torch.max(outputs, 1)

    return class_names[preds.item()]